/**
# About how to add memory tracker in google benchmark

- SEE: https://github.com/google/benchmark/issues/1217
- SEE ALSO:
https://github.com/google/benchmark/blob/main/docs/user_guide.md#custom-counters

# Use tcmalloc to track memory

- SEE: https://github.com/google/tcmalloc/issues/44
- SEE ALSO for the TCMalloc Implementation:
https://github.com/google/tcmalloc/commit/61b60f4f796c11dfd0ec5dbb8a8e1aaf188b6fbd

# Write custom implementation

- REFER TO: https://github.com/archibate/mallocvis/blob/main/malloc_hook.cpp
*/

#pragma once

#include <benchmark/benchmark.h>
#include <gperftools/malloc_hook.h>

#include <algorithm>
#include <tuple>

class MemTracker {
   public:
    static MemTracker& getInstance() {
        static MemTracker instance;
        return instance;
    }

    void AddGNumNew(benchmark::IterationCount num) { g_num_new += num; }
    void AddGSumSizeNew(benchmark::IterationCount size) {
        g_sum_size_new += size;
    }
    // get g_num_new
    benchmark::IterationCount GetGNumNew() const { return g_num_new; }
    // get g_sum_size_new
    benchmark::IterationCount GetGSumSizeNew() const { return g_sum_size_new; }

    // 删除拷贝和赋值操作
    MemTracker(const MemTracker&) = delete;
    MemTracker& operator=(const MemTracker&) = delete;

   private:
    benchmark::IterationCount g_num_new;
    benchmark::IterationCount g_sum_size_new;
    MemTracker() : g_num_new(0), g_sum_size_new(0){};
    ~MemTracker() = default;
};

auto new_hook = [](const void*, size_t size) {
    MemTracker::getInstance().AddGNumNew(1);
    MemTracker::getInstance().AddGSumSizeNew(size);
};

#define BEFORE_TEST                                 \
    benchmark::IterationCount num_new =             \
        MemTracker::getInstance().GetGNumNew();     \
    benchmark::IterationCount sum_size_new =        \
        MemTracker::getInstance().GetGSumSizeNew(); \
    MallocHook::AddNewHook(new_hook);

#define AFTER_TEST                                                        \
    MallocHook::RemoveNewHook(new_hook);                                  \
    auto iter = state.iterations();                                       \
    auto g_num_new = MemTracker::getInstance().GetGNumNew();              \
    auto g_sum_size_new = MemTracker::getInstance().GetGSumSizeNew();     \
    state.counters["#new"] = (g_num_new - num_new) / iter;                \
    state.counters["sum_new_B"] = (g_sum_size_new - sum_size_new) / iter; \
    state.counters["avg_new_B"] =                                         \
        (g_sum_size_new - sum_size_new) / (g_num_new - num_new);